package com.yunlongn.async.handle;
import com.yunlongn.async.content.NameThreadFactory;
import com.yunlongn.async.exception.ExceptionHandler;
import com.yunlongn.async.function.AsyncRunnableAction;
import com.yunlongn.async.thread.ShardingRunnable;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.util.CollectionUtils;
import org.springframework.util.ObjectUtils;

import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

/**
 * 分片异步处理器
 * @author yunlongn
 * @since 20210907
 */
public class AsyncShardingHandle {
    private final static Logger log = LoggerFactory.getLogger(AsyncShardingHandle.class);
    /**
     * 分片处理的线程池
     */
    private final ExecutorService executorService;
    /**
     * 执行线程数 用于本身的异步。
     */
    private final static ExecutorService RUNNABLE_EXECUTOR_SERVICE;
    private Semaphore semaphore;
    private CountDownLatch latch;
    private int maximumPoolSize = 0;
    /**
     * 分片任务 传入的List内容包装此处
     */
    private List<ShardingRunnable> shardingRunnableList;
    /**
     * 执行线程池
     */
    private ExceptionHandler<Object> exceptionHandler;

    static {
        RUNNABLE_EXECUTOR_SERVICE = new ThreadPoolExecutor(
                2, 8,
                60, TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(),
                new NameThreadFactory("ASH-EXECUTOR-", true));
        Runtime.getRuntime().addShutdownHook(new Thread(AsyncShardingHandle::destroy));
    }
    public AsyncShardingHandle(ExecutorService executorService) {
        if (executorService instanceof  ThreadPoolExecutor) {
            ThreadPoolExecutor threadPoolExecutor = (ThreadPoolExecutor)executorService;
            this.maximumPoolSize = threadPoolExecutor.getMaximumPoolSize();
            if (Executors.defaultThreadFactory().equals(threadPoolExecutor.getThreadFactory())) {
                threadPoolExecutor.setThreadFactory(new NameThreadFactory("ASH-Thread-", true));
            }
        }
        this.executorService = executorService;
        this.semaphore = new Semaphore(3);
        this.shardingRunnableList = new ArrayList<>();
    }

    private static void destroy() {
        if (RUNNABLE_EXECUTOR_SERVICE != null && !RUNNABLE_EXECUTOR_SERVICE.isShutdown()) {
            RUNNABLE_EXECUTOR_SERVICE.shutdownNow();
        }
    }

    public AsyncShardingHandle exceptionHandler(ExceptionHandler<Object> exceptionHandler) {
        this.exceptionHandler = exceptionHandler;
        return this;
    }

    /**
     * 设置执行线程数
     * @param threads 线程数
     */
    public AsyncShardingHandle threadsNum(int threads) {
        if (threads <= 1) {
            throw new RuntimeException("threads num error 请设置至少两个线程来执行异步任务");
        }
        if (maximumPoolSize < threads) {
            throw new RuntimeException("设置线程数请不要超过线程池最大线程数");
        }
        this.semaphore = new Semaphore(threads);
        return this;
    }

    public <T> AsyncShardingHandle toRunnable(Collection<T> data, AsyncRunnableAction<T> action) {
        if (CollectionUtils.isEmpty(data)) {
            this.shardingRunnableList = new ArrayList<>(0);
            return this;
        }
        Map<String, String> mdcContextMap = MDC.getCopyOfContextMap();
        this.shardingRunnableList = data.parallelStream()
        .filter(Objects::nonNull)
        .map(item -> new ShardingRunnable(item) {
            @Override
            public void run() {
                try {
                    // 日志变量透传
                    if (!ObjectUtils.isEmpty(mdcContextMap)) {
                        MDC.setContextMap(mdcContextMap);
                    }
                    action.action(item);
                } finally {
                    MDC.clear();
                }
            }
        }).collect(Collectors.toList());
        return this;
    }


    /**
     * 执行所有任务 并等待任务结束
     */
    public void executes() throws InterruptedException {
        executesAsync();
        await();
    }

    public void await() throws InterruptedException {
        if (latch != null) {
            latch.await();
        }
    }

    /**
     * 按照进程数执行任务
     */
    public void executesAsync() {
        if (shardingRunnableList.isEmpty()) {
            return;
        }
        // 统计完成数
        latch = new CountDownLatch(shardingRunnableList.size());
        // 使用线程执行。不阻塞主线程
        RUNNABLE_EXECUTOR_SERVICE.submit(() ->
                shardingRunnableList.forEach(item -> {
            try {
                // 信号量 +1
                semaphore.acquire();
            } catch (InterruptedException e) {
                log.error("semaphore.acquire e ： {}", ExceptionUtils.getStackTrace(e));
            }
            executorService.submit(() -> {
                try {
                    item.run();
                } catch (Exception e) {
                    if (exceptionHandler != null) {
                        exceptionHandler.onException(item.getData(), e);
                    } else {
                        log.error("runnable e ={}", ExceptionUtils.getStackTrace(e));
                    }
                } finally {
                    // 信号量 -1
                    semaphore.release();
                    // 剩余任务 -1
                    latch.countDown();
                }
            });
        }));
    }

}
